#!/usr/bin/python
# pylint: disable-msg = E1101, C0301
'''
Created by Colin Gillespie <c.gillespie@ncl.ac.uk>
Last Changed on: $LastChangedDate: 2008-03-13 19:50:58 +0000 (Thu, 13 Mar 2008) $
Rev: $Rev: 73 $
Main part of the algorthim

'''

__revision__ = '$Rev: 73 $'

#Import non-standard modules
from mrange import mrange
from Factorial import binomial
from VariableDict import Variables
from swiginac import numeric, expand

def getCumulantEquation(cum_equ, reaction_list):
    assert len(cum_equ) == len(reaction_list[0].species_order)
    assert sum(cum_equ) <= 2

    if sum(cum_equ) == 1:
        return getCumulantEquation1(cum_equ, reaction_list)
    elif sum(cum_equ) == 2 and max(cum_equ) == 2:
        cum_equ = list(cum_equ)
        index = cum_equ.index(2)
        mu2 = getCumulantEquation1(tuple(cum_equ), reaction_list)
        cum_equ[index] = 1
        mu1 = 2*getCumulantEquation1(tuple(cum_equ), reaction_list)*Variables().getCumulant(tuple(cum_equ))
        return expand(mu2 - mu1)
    else:
        cum_equ = list(cum_equ)
        mu11 = getCumulantEquation1(tuple(cum_equ), reaction_list)

        index1 = cum_equ.index(1)
        cum_equ[index1] = 0
        mu01 = getCumulantEquation1(tuple(cum_equ), reaction_list)
        kappa01 = Variables().getCumulant(tuple(cum_equ))
        index2 = cum_equ.index(1)
        cum_equ[index1] = 1
        cum_equ[index2] = 0
       
        mu10 = getCumulantEquation1(tuple(cum_equ), reaction_list)
        kappa10 = Variables().getCumulant(tuple(cum_equ))
        return expand(mu11 - mu01*kappa10 - mu10*kappa01)

def getCumulantEquation1(cum_equ, reaction_list):

    entry = 0
    for reaction in reaction_list:
        for coeffs, values in reaction.rate.items():
            entry += values*inner_sum_cum(cum_equ, reaction, coeffs)
    return entry.expand()                      

def inner_sum_cum(mom_equ, reaction, coeffs):
    entry = 0
    specific_entry = 0
    for i2_tup in mrange(mom_equ):
        if sum(i2_tup) != 0:
            specific_entry = numeric(1)
            for j, value in enumerate(i2_tup):
                specific_entry *= binomial(mom_equ[j], value)*pow(reaction.stoic[j], value)
            i_j_diff = tuple([n1-k1+j1 for n1, k1, j1 in zip(mom_equ, i2_tup, coeffs)])
            specific_entry *= Variables().getCumulant(i_j_diff)

        entry += specific_entry
            
    return entry

def getMomentEquation(mom_equ, reaction_list):
    '''
        Input: mom_equ - a tuple indicating which moment equation you wish
        reaction_list: a list of reactions with rate latws
        output: cumulant equation
    '''
    
    assert len(mom_equ) == len(reaction_list[0].species_order)
    entry = 0
    for reaction in reaction_list:
        for coeffs, values in reaction.rate.items():
            
            entry += values*inner_sum(mom_equ, reaction, coeffs)
            #print values, inner_sum(mom_equ, reaction, coeffs)
    return entry.expand()                      
            
            
            
def inner_sum(mom_equ, reaction, coeffs):
    entry = 0
    specific_entry = 0
    for i2_tup in mrange(mom_equ):
        if sum(i2_tup) != 0:
            specific_entry = numeric(1)
            for j, value in enumerate(i2_tup):
                specific_entry *= binomial(mom_equ[j], value)*pow(reaction.stoic[j], value)
            i_j_diff = tuple([n1-k1+j1 for n1, k1, j1 in zip(mom_equ, i2_tup, coeffs)])
            #print i_j_diff, binomial(mom_equ[j],value)*pow(reaction.stoic[j],value)
            specific_entry *= Variables().getMoment(i_j_diff)

        entry += specific_entry
            
    return entry

                



if __name__ == '__main__':
    Variables().initialise()
    import GetReactionInformation
    #gdd = GetReactionInformation.GetStoichiometryDict(sbml_file = 'xml_files/sis.xml')
    gdd = GetReactionInformation.GetStoichiometryDict(sbml_file = 'xml_files/Hsp90model1.xml')    
    react_list = gdd.getReactionList()
    
    print getCumulantEquation((1,0,0,0,0,0,0,0,0,0,0,0,0,0,1), react_list)
#    print getMomentEquation((1,0,0,0,0,0,0,0,0,0,0,0,0,0,1), reaction_list)
    #print getMomentEquation((2,), reaction_list)
#    print getMomentEquation((2,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0), reaction_list)
    
#    for value in reaction_list:
#        print value.sbml_id, value.species_order, value.stoic
#        print value.rate
